#include "kernel_operator.h"
#include <cstdint>
#include <cmath>

constexpr uint32_t ISCOMPLEXOBJ_BFLOAT16 = 0;            // 对应 DT_BF16，bf16 类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT16 = 1;             // 对应 DT_FLOAT16，half（float16）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT32 = 2;             // 对应 DT_FLOAT，float32 类型
constexpr uint32_t ISCOMPLEXOBJ_INT8 = 3;                // 对应 DT_INT8，int8 类型
constexpr uint32_t ISCOMPLEXOBJ_INT16 = 4;               // 对应 DT_INT16，int16 类型
constexpr uint32_t ISCOMPLEXOBJ_INT32 = 5;               // 对应 DT_INT32，int32 类型
constexpr uint32_t ISCOMPLEXOBJ_UINT8 = 6;               // 对应 DT_UINT8，u8（无符号int8）类型
constexpr uint32_t ISCOMPLEXOBJ_UINT16 = 7;              // 对应 DT_UINT16，u16（无符号int16）类型
constexpr uint32_t ISCOMPLEXOBJ_UINT32 = 8;              // 对应 DT_UINT32，u32（无符号int32）类型
constexpr uint32_t ISCOMPLEXOBJ_INT64 = 9;               // 对应 DT_INT64，int64 类型
constexpr uint32_t ISCOMPLEXOBJ_UINT64 = 10;             // 对应 DT_UINT64，u64（无符号int64）类型
constexpr uint32_t ISCOMPLEXOBJ_BOOL = 11;               // 对应 DT_BOOL，bool 类型
constexpr uint32_t ISCOMPLEXOBJ_DOUBLE = 12;             // 对应 DT_DOUBLE，double（float64）类型
constexpr uint32_t ISCOMPLEXOBJ_STRING = 13;             // 对应 DT_STRING，string（字符串）类型
constexpr uint32_t ISCOMPLEXOBJ_DUAL_SUB_INT8 = 14;      // 对应 DT_DUAL_SUB_INT8，dual output int8（双输出int8）类型
constexpr uint32_t ISCOMPLEXOBJ_DUAL_SUB_UINT8 = 15;     // 对应 DT_DUAL_SUB_UINT8，dual output u8（双输出无符号int8）类型
constexpr uint32_t ISCOMPLEXOBJ_COMPLEX64 = 16;          // 对应 DT_COMPLEX64，complex64（复数64位）类型
constexpr uint32_t ISCOMPLEXOBJ_COMPLEX128 = 17;         // 对应 DT_COMPLEX128，complex128（复数128位）类型
constexpr uint32_t ISCOMPLEXOBJ_QINT8 = 18;              // 对应 DT_QINT8，qint8（量化int8）类型
constexpr uint32_t ISCOMPLEXOBJ_QINT16 = 19;             // 对应 DT_QINT16，qint16（量化int16）类型
constexpr uint32_t ISCOMPLEXOBJ_QINT32 = 20;             // 对应 DT_QINT32，qint32（量化int32）类型
constexpr uint32_t ISCOMPLEXOBJ_QUINT8 = 21;             // 对应 DT_QUINT8，quint8（量化无符号int8）类型
constexpr uint32_t ISCOMPLEXOBJ_QUINT16 = 22;            // 对应 DT_QUINT16，quint16（量化无符号int16）类型
constexpr uint32_t ISCOMPLEXOBJ_RESOURCE = 23;           // 对应 DT_RESOURCE，resource（资源）类型
constexpr uint32_t ISCOMPLEXOBJ_STRING_REF = 24;         // 对应 DT_STRING_REF，string ref（字符串引用）类型
constexpr uint32_t ISCOMPLEXOBJ_DUAL = 25;               // 对应 DT_DUAL，dual output（双输出）类型
constexpr uint32_t ISCOMPLEXOBJ_VARIANT = 26;            // 对应 DT_VARIANT，dt_variant（变体）类型
constexpr uint32_t ISCOMPLEXOBJ_UNDEFINED = 27;          // 对应 DT_UNDEFINED，未设置的DataType字段
constexpr uint32_t ISCOMPLEXOBJ_INT4 = 28;               // 对应 DT_INT4，int4 类型
constexpr uint32_t ISCOMPLEXOBJ_UINT1 = 29;              // 对应 DT_UINT1，u1（无符号int1）类型
constexpr uint32_t ISCOMPLEXOBJ_INT2 = 30;               // 对应 DT_INT2，int2 类型
constexpr uint32_t ISCOMPLEXOBJ_UINT2 = 31;              // 对应 DT_UINT2，u2（无符号int2）类型
constexpr uint32_t ISCOMPLEXOBJ_COMPLEX32 = 32;          // 对应 DT_COMPLEX32，complex32（复数32位）类型
constexpr uint32_t ISCOMPLEXOBJ_HIFLOAT8 = 33;           // 对应 DT_HIFLOAT8，hif8（高精度float8）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT8_E5M2 = 34;        // 对应 DT_FLOAT8_E5M2，fp8_e5m2（float8格式，5位指数2位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT8_E4M3FN = 35;      // 对应 DT_FLOAT8_E4M3FN，fp8_e4m3（float8格式，4位指数3位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT8_E8M0 = 36;        // 对应 DT_FLOAT8_E8M0，fp8_e8m0（float8格式，8位指数0位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT6_E3M2 = 37;        // 对应 DT_FLOAT6_E3M2，fp6_e3m2（float6格式，3位指数2位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT6_E2M3 = 38;        // 对应 DT_FLOAT6_E2M3，fp6_e2m3（float6格式，2位指数3位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT4_E2M1 = 39;        // 对应 DT_FLOAT4_E2M1，fp4_e2m1（float4格式，2位指数1位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_FLOAT4_E1M2 = 40;        // 对应 DT_FLOAT4_E1M2，fp4_e1m2（float4格式，1位指数2位尾数）类型
constexpr uint32_t ISCOMPLEXOBJ_MAX = 41;                // 对应 DT_MAX，标记AscendC数据类型的边界

constexpr int32_t TOTAL_LENGTH = 8 * 2048;
constexpr int32_t USE_CORE_NUM = 8;
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM; 
constexpr int32_t TILE_NUM = 8;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM;

// 定义复数类型
struct complex64 {
    float real;
    float imag;
};

class KernelIsComplexObj {
public:
    __aicore__ inline KernelIsComplexObj() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR out) {
        this->dataTypeId = 16; // 假设输入数据类型为 complex64，可以根据实际需求进行修改
        int32_t blockIdx = AscendC::GetBlockIdx();
        xGm.SetGlobalBuffer((__gm__ complex64*)x + blockIdx * BLOCK_LENGTH, BLOCK_LENGTH);
        outGm.SetGlobalBuffer((__gm__ int8_t*)out, 1);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, TILE_LENGTH * sizeof(complex64));
        pipe.InitBuffer(outQueueBool, BUFFER_NUM, 1);
    }

    __aicore__ inline void Process() {
        CopyIn(0);
        Compute(0);
        CopyOut(0);
    }

private:
    __aicore__ inline void CopyIn(int32_t progress) {
        int32_t tileOffset = progress * TILE_LENGTH;
        if (tileOffset >= BLOCK_LENGTH) return;

        AscendC::LocalTensor<complex64> xLocal = inQueueX.AllocTensor<complex64>();

        AscendC::DataCopy(xLocal, xGm[progress * TILE_LENGTH], TILE_LENGTH);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress) {
        AscendC::LocalTensor<complex64> xLocal = inQueueX.DeQue<complex64>();
        AscendC::LocalTensor<int8_t> outLocal = outQueueBool.AllocTensor<int8_t>();
        
        bool isComplex = (dataTypeId == ISCOMPLEXOBJ_COMPLEX64 || dataTypeId == ISCOMPLEXOBJ_COMPLEX128);

        if (!isComplex) {
            outLocal.SetValue(0, 0);
        }
        else{
            outLocal.SetValue(0, 1);
        }

        outQueueBool.EnQue(outLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress) {
        AscendC::LocalTensor<int8_t> outLocal = outQueueBool.DeQue<int8_t>();
        AscendC::DataCopyExtParams copyParams{1, (uint32_t)(1), 0, 0, 0};
        AscendC::DataCopyPad(outGm[0], outLocal, copyParams);
        
        // AscendC::DataCopy(outGm[0], outLocal, 1);
        outQueueBool.FreeTensor(outLocal);
    }
    
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueBool;
    AscendC::GlobalTensor<complex64> xGm;
    AscendC::GlobalTensor<int8_t> outGm;
    AscendC::DataCopyExtParams copyParams;

    uint32_t dataTypeId;
};


extern "C" __global__ __aicore__ void is_complex_obj_custom(GM_ADDR x, GM_ADDR out)
{
    KernelIsComplexObj op;
    op.Init(x, out);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void is_complex_obj_custom_do(uint32_t blockDim, void *stream, void *x, void *out)
{
    is_complex_obj_custom<<<blockDim, nullptr, stream>>>(x, out);
}
#endif